{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/facebookresearch/habitat-sim/blob/main/examples/tutorials/colabs/ReplicaCAD_quickstart.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#Habitat-sim ReplicaCAD Quickstart\n",
    "\n",
    "This brief Colab tutorial demonstrates loading the ReplicaCAD dataset in Habitat-sim from a SceneDataset and rendering a short video of agent navigation with physics simulation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Installation { display-mode: \"form\" }\n",
    "# @markdown (double click to show code).\n",
    "\n",
    "!curl -L https://raw.githubusercontent.com/facebookresearch/habitat-sim/main/examples/colab_utils/colab_install.sh | NIGHTLY=true bash -s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# @title Path Setup and Imports { display-mode: \"form\" }\n",
    "# @markdown (double click to show code).\n",
    "\n",
    "%cd /content/habitat-sim\n",
    "## [setup]\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import git\n",
    "import magnum as mn\n",
    "\n",
    "import habitat_sim\n",
    "from habitat_sim.utils import viz_utils as vut\n",
    "\n",
    "try:\n",
    "    import ipywidgets as widgets\n",
    "    from IPython.display import display as ipydisplay\n",
    "\n",
    "    # For using jupyter/ipywidget IO components\n",
    "\n",
    "    HAS_WIDGETS = True\n",
    "except ImportError:\n",
    "    HAS_WIDGETS = False\n",
    "\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    os.environ[\"IMAGEIO_FFMPEG_EXE\"] = \"/usr/bin/ffmpeg\"\n",
    "\n",
    "repo = git.Repo(\".\", search_parent_directories=True)\n",
    "dir_path = repo.working_tree_dir\n",
    "%cd $dir_path\n",
    "data_path = os.path.join(dir_path, \"data\")\n",
    "# fmt: off\n",
    "output_directory = \"examples/tutorials/replica_cad_output/\"  # @param {type:\"string\"}\n",
    "# fmt: on\n",
    "output_path = os.path.join(dir_path, output_directory)\n",
    "if not os.path.exists(output_path):\n",
    "    os.mkdir(output_path)\n",
    "\n",
    "# define some globals the first time we run.\n",
    "if \"sim\" not in globals():\n",
    "    global sim\n",
    "    sim = None\n",
    "    global obj_attr_mgr\n",
    "    obj_attr_mgr = None\n",
    "    global stage_attr_mgr\n",
    "    stage_attr_mgr = None\n",
    "    global rigid_obj_mgr\n",
    "    rigid_obj_mgr = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# @title Define Configuration Utility Functions { display-mode: \"form\" }\n",
    "# @markdown (double click to show code)\n",
    "\n",
    "# @markdown This cell defines a number of utility functions used throughout the tutorial to make simulator reconstruction easy:\n",
    "# @markdown - make_cfg\n",
    "# @markdown - make_default_settings\n",
    "# @markdown - make_simulator_from_settings\n",
    "\n",
    "\n",
    "def make_cfg(settings):\n",
    "    sim_cfg = habitat_sim.SimulatorConfiguration()\n",
    "    sim_cfg.gpu_device_id = 0\n",
    "    sim_cfg.scene_dataset_config_file = settings[\"scene_dataset\"]\n",
    "    sim_cfg.scene_id = settings[\"scene\"]\n",
    "    sim_cfg.enable_physics = settings[\"enable_physics\"]\n",
    "    # Specify the location of the scene dataset\n",
    "    if \"scene_dataset_config\" in settings:\n",
    "        sim_cfg.scene_dataset_config_file = settings[\"scene_dataset_config\"]\n",
    "    if \"override_scene_light_defaults\" in settings:\n",
    "        sim_cfg.override_scene_light_defaults = settings[\n",
    "            \"override_scene_light_defaults\"\n",
    "        ]\n",
    "    if \"scene_light_setup\" in settings:\n",
    "        sim_cfg.scene_light_setup = settings[\"scene_light_setup\"]\n",
    "\n",
    "    # Note: all sensors must have the same resolution\n",
    "    sensor_specs = []\n",
    "    color_sensor_1st_person_spec = habitat_sim.CameraSensorSpec()\n",
    "    color_sensor_1st_person_spec.uuid = \"color_sensor_1st_person\"\n",
    "    color_sensor_1st_person_spec.sensor_type = habitat_sim.SensorType.COLOR\n",
    "    color_sensor_1st_person_spec.resolution = [\n",
    "        settings[\"height\"],\n",
    "        settings[\"width\"],\n",
    "    ]\n",
    "    color_sensor_1st_person_spec.position = [0.0, settings[\"sensor_height\"], 0.0]\n",
    "    color_sensor_1st_person_spec.orientation = [\n",
    "        settings[\"sensor_pitch\"],\n",
    "        0.0,\n",
    "        0.0,\n",
    "    ]\n",
    "    color_sensor_1st_person_spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE\n",
    "    sensor_specs.append(color_sensor_1st_person_spec)\n",
    "\n",
    "    # Here you can specify the amount of displacement in a forward action and the turn angle\n",
    "    agent_cfg = habitat_sim.agent.AgentConfiguration()\n",
    "    agent_cfg.sensor_specifications = sensor_specs\n",
    "\n",
    "    return habitat_sim.Configuration(sim_cfg, [agent_cfg])\n",
    "\n",
    "\n",
    "def make_default_settings():\n",
    "    settings = {\n",
    "        \"width\": 1280,  # Spatial resolution of the observations\n",
    "        \"height\": 720,\n",
    "        \"scene_dataset\": \"data/replica_cad/replicaCAD.scene_dataset_config.json\",  # dataset path\n",
    "        \"scene\": \"NONE\",  # Scene path\n",
    "        \"default_agent\": 0,\n",
    "        \"sensor_height\": 1.5,  # Height of sensors in meters\n",
    "        \"sensor_pitch\": 0,  # sensor pitch (x rotation in rads)\n",
    "        \"seed\": 1,\n",
    "        \"enable_physics\": True,  # enable dynamics simulation\n",
    "    }\n",
    "    return settings\n",
    "\n",
    "\n",
    "def make_simulator_from_settings(sim_settings):\n",
    "    cfg = make_cfg(sim_settings)\n",
    "    # clean-up the current simulator instance if it exists\n",
    "    global sim\n",
    "    global obj_attr_mgr\n",
    "    global prim_attr_mgr\n",
    "    global stage_attr_mgr\n",
    "    global rigid_obj_mgr\n",
    "    global metadata_mediator\n",
    "\n",
    "    if sim != None:\n",
    "        sim.close()\n",
    "    # initialize the simulator\n",
    "    sim = habitat_sim.Simulator(cfg)\n",
    "    # Managers of various Attributes templates\n",
    "    obj_attr_mgr = sim.get_object_template_manager()\n",
    "    obj_attr_mgr.load_configs(str(os.path.join(data_path, \"objects/example_objects\")))\n",
    "    prim_attr_mgr = sim.get_asset_template_manager()\n",
    "    stage_attr_mgr = sim.get_stage_template_manager()\n",
    "    # Manager providing access to rigid objects\n",
    "    rigid_obj_mgr = sim.get_rigid_object_manager()\n",
    "    # get metadata_mediator\n",
    "    metadata_mediator = sim.metadata_mediator\n",
    "\n",
    "    # UI-populated handles used in various cells.  Need to initialize to valid\n",
    "    # value in case IPyWidgets are not available.\n",
    "    # Holds the user's desired scene handle\n",
    "    global selected_scene\n",
    "    selected_scene = \"NONE\"\n",
    "\n",
    "\n",
    "# [/setup]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Define Simulation Utility Function { display-mode: \"form\" }\n",
    "# @markdown (double click to show code)\n",
    "def simulate(sim, dt=1.0, get_frames=True):\n",
    "    # simulate dt seconds at 60Hz to the nearest fixed timestep\n",
    "    print(\"Simulating {:.3f} world seconds.\".format(dt))\n",
    "    observations = []\n",
    "    start_time = sim.get_world_time()\n",
    "    while sim.get_world_time() < start_time + dt:\n",
    "        sim.step_physics(1.0 / 60.0)\n",
    "        if get_frames:\n",
    "            observations.append(sim.get_sensor_observations())\n",
    "    return observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Define Colab GUI Utility Functions { display-mode: \"form\" }\n",
    "# @markdown (double click to show code)\n",
    "\n",
    "# @markdown This cell provides utility functions to build and manage IPyWidget interactive components.\n",
    "\n",
    "# Event handler for dropdowns displaying file-based object handles\n",
    "def on_scene_ddl_change(ddl_values):\n",
    "    global selected_scene\n",
    "    selected_scene = ddl_values[\"new\"]\n",
    "    return selected_scene\n",
    "\n",
    "\n",
    "# Build a dropdown list holding obj_handles and set its event handler\n",
    "def set_handle_ddl_widget(scene_handles, sel_handle, on_change):\n",
    "    sel_handle = scene_handles[0]\n",
    "    descStr = \"Available Scenes:\"\n",
    "    style = {\"description_width\": \"300px\"}\n",
    "    obj_ddl = widgets.Dropdown(\n",
    "        options=scene_handles,\n",
    "        value=sel_handle,\n",
    "        description=descStr,\n",
    "        style=style,\n",
    "        disabled=False,\n",
    "        layout={\"width\": \"max-content\"},\n",
    "    )\n",
    "\n",
    "    obj_ddl.observe(on_change, names=\"value\")\n",
    "    return obj_ddl, sel_handle\n",
    "\n",
    "\n",
    "def set_button_launcher(desc):\n",
    "    button = widgets.Button(\n",
    "        description=desc,\n",
    "        layout={\"width\": \"max-content\"},\n",
    "    )\n",
    "    return button\n",
    "\n",
    "\n",
    "# Builds widget-based UI components\n",
    "def build_widget_ui(metadata_mediator):\n",
    "    # Holds the user's desired scene\n",
    "    global selected_scene\n",
    "    selected_scene = \"NONE\"\n",
    "\n",
    "    # Construct DDLs and assign event handlers\n",
    "    # All file-based object template handles\n",
    "    scene_handles = metadata_mediator.get_scene_handles()\n",
    "    # If not using widgets, set as first available handle\n",
    "    if not HAS_WIDGETS:\n",
    "        selected_scene = scene_handles[0]\n",
    "        return\n",
    "\n",
    "    # Build widgets\n",
    "    scene_obj_ddl, selected_scene = set_handle_ddl_widget(\n",
    "        scene_handles,\n",
    "        selected_scene,\n",
    "        on_scene_ddl_change,\n",
    "    )\n",
    "\n",
    "    # Display DDLs\n",
    "    ipydisplay(scene_obj_ddl)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import argparse\n",
    "\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--no-display\", dest=\"display\", action=\"store_false\")\n",
    "    parser.add_argument(\"--no-make-video\", dest=\"make_video\", action=\"store_false\")\n",
    "    parser.set_defaults(show_video=True, make_video=True)\n",
    "    args, _ = parser.parse_known_args()\n",
    "    show_video = args.display\n",
    "    display = args.display\n",
    "    make_video = args.make_video\n",
    "else:\n",
    "    show_video = False\n",
    "    make_video = False\n",
    "    display = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# View ReplicaCAD in Habitat-sim\n",
    "Use the code in this section to view assets in the Habitat-sim engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [initialize]\n",
    "# @title Initialize Simulator{ display-mode: \"form\" }\n",
    "\n",
    "sim_settings = make_default_settings()\n",
    "make_simulator_from_settings(sim_settings)\n",
    "# [/initialize]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Select a SceneInstance: { display-mode: \"form\" }\n",
    "# @markdown Select a scene from the dropdown and then run the next cell to load and simulate that scene and produce a visualization of the result.\n",
    "\n",
    "build_widget_ui(sim.metadata_mediator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Select Scene and Simulate!\n",
    "This cell will load the scene selected above, simulate, and produce a visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global selected_scene\n",
    "if sim_settings[\"scene\"] != selected_scene:\n",
    "    sim_settings[\"scene\"] = selected_scene\n",
    "    make_simulator_from_settings(sim_settings)\n",
    "\n",
    "observations = []\n",
    "start_time = sim.get_world_time()\n",
    "while sim.get_world_time() < start_time + 4.0:\n",
    "    sim.agents[0].scene_node.rotate(mn.Rad(mn.math.pi_half / 60.0), mn.Vector3(0, 1, 0))\n",
    "    sim.step_physics(1.0 / 60.0)\n",
    "    if make_video:\n",
    "        observations.append(sim.get_sensor_observations())\n",
    "\n",
    "# video rendering of carousel view\n",
    "video_prefix = \"ReplicaCAD_scene_view\"\n",
    "if make_video:\n",
    "    vut.make_video(\n",
    "        observations,\n",
    "        \"color_sensor_1st_person\",\n",
    "        \"color\",\n",
    "        output_path + video_prefix,\n",
    "        open_vid=show_video,\n",
    "        video_dims=[1280, 720],\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Habitat-sim ReplicaCAD Quickstart",
   "private_outputs": true,
   "provenance": []
  },
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "nb_python//py:percent,colabs//ipynb",
   "main_language": "python",
   "notebook_metadata_filter": "all"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
